/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;

use anyhow::Error;
use async_trait::async_trait;
use bookmarks::BookmarkTransactionError;
use context::CoreContext;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::Timestamp;
use sql_ext::Transaction;

pub type RebasedChangesets = HashMap<ChangesetId, (ChangesetId, Timestamp)>;

#[async_trait]
pub trait PushrebaseHook: Send + Sync + 'static {
    /// in_critical_section is called once per pushrebase attempt, right after loading the target
    /// bookmark for pushrebase. It should return a PushrebaseCommitHook for further processing of
    /// the changeset generated by pushrebase.
    /// It can assume no more commits will land to this bookmark until the rest of the pushrebase completes.
    async fn in_critical_section(
        &self,
        _ctx: &CoreContext,
        _old_bookmark_value: Option<ChangesetId>,
    ) -> Result<Box<dyn PushrebaseCommitHook>, Error>;
}

#[async_trait]
pub trait PushrebaseCommitHook: Send + Sync + 'static {
    /// post_rebase_changeset is called once per pushrebased changeset, with the ID of the
    /// changeset that was rebased, and a mutable reference to the resulting changeset. This can be
    /// used to modify the pushrebased commit.
    fn post_rebase_changeset(
        &mut self,
        _bcs_old: ChangesetId,
        _bcs_new: &mut BonsaiChangesetMut,
    ) -> Result<(), Error> {
        Ok(())
    }

    /// into_transaction_hook is called after all commits have been pushrebased and we're going to
    /// attempt to move the bookmark to pushrebase onto. This should capture any data necessary and
    /// return a PushrebaseTransactionHook.
    async fn into_transaction_hook(
        self: Box<Self>,
        _ctx: &CoreContext,
        changesets: &RebasedChangesets,
    ) -> Result<Box<dyn PushrebaseTransactionHook>, Error>;
}

#[async_trait]
pub trait PushrebaseTransactionHook: Send + Sync + 'static {
    /// populate_transaction is called once per attempt to update bookmarks in the DB (the
    /// operation might be retried if the write for reasons that are NOT another push having won
    /// the race). This can add extra writes to the transaction if desired.
    async fn populate_transaction(
        &self,
        ctx: &CoreContext,
        txn: Transaction,
    ) -> Result<Transaction, BookmarkTransactionError>;
}
